"""
Name:        lwps - Lightweight Python Stack
Version:     0.1
Author:      Itamar Tal
"""

###############################################################################
### Imports
###############################################################################

import sys
import imp
import time
import dpkt
import logging
import threading

from ip import IP
from arp import ARP
from udp import UDP, UDP_PROTOCOL
from tcp import TCP, TCP_PROTOCOL
from ethernet import Ethernet, ETHERNET_BROADCAST, ETHERNET_TYPE_IP, ETHERNET_TYPE_ARP

# import pywinpcap
pywinpcap = imp.load_dynamic("pywinpcap", r"d:\Development\PyWinPCap\Debug\pywinpcap.pyd")
sys.modules["pywinpcap"] = pywinpcap

logging.basicConfig(format="%(asctime)s: %(levelname)s: %(module)s:\t%(message)s",
                    datefmt="%H:%M:%S", # "%Y-%m-%d %H:%M:%S",
                    level=logging.DEBUG)

###############################################################################
### Consts
###############################################################################

SOURCE_MAC_ADDR = "00:06:4f:58:99:40"
SOURCE_IP_ADDR = "192.168.1.123"

###############################################################################
### Script
###############################################################################

class RecvPacketThread(threading.Thread):
    def __init__(self, network):
        threading.Thread.__init__(self)
        
        self.m_network = network
        self.m_keep_running = True
        
        self.m_handlers = {}
    
    def _pkt_handler(self, device_obj, packet_obj, context):
        logging.debug("Packet Received: %s" % (str(packet_obj), ))
        
        done_handlers = []
        
        # call all handlers
        for handler in self.m_handlers.keys():
            dpkt_pkt = dpkt.ethernet.Ethernet(packet_obj.data)
            
            # check if packet pass filter
            if self.m_handlers[handler][0](dpkt_pkt):
                # call handler, check if we're done
                if not handler(device_obj, dpkt_pkt):
                    done_handlers.append(handler)
        
        # remove done handlers (and release the lock)
        for done_handler in done_handlers:
            self.m_handlers.pop(done_handler)[1].release()
        
        return self.m_keep_running
    
    def add_handler(self, filter_callback, handler_callback):
        """add_handler(handler) -> Lock
        
        Add a new packet handler. Packet handler is a callback function
        called upon each received packet. If the handler is done, it
        should return False. The function returns a lock, which is
        locked until the handler is done.
        """
        
        self.m_handlers[handler_callback] = (filter_callback, threading.Lock())
        self.m_handlers[handler_callback][1].acquire()
        return self.m_handlers[handler_callback][1]
    
    def kill(self):
        logging.debug("Killing packet receiving thread...")
        self.m_keep_running = False
    
    def run(self):
        logging.info("Waiting for packets...")
        
        try:
            self.m_network.m_capture.loop(self._pkt_handler, self)
        except pywinpcap.error, e:
            logging.debug("Capturing loop broke out: %s" % (e, ))

class Network:
    def __init__(self, capture, src_ip=SOURCE_IP_ADDR, src_mac=SOURCE_MAC_ADDR):
        self.m_capture = capture
        
        # set local network parameters
        self.m_src_mac = Ethernet.mac_addr_aton(src_mac)
        self.m_src_ip = IP.ip_addr_aton(src_ip)
        
        
        # start reading packets
        self.m_recv_pkt_thread = RecvPacketThread(self)
        self.m_recv_pkt_thread.start()
        
        # set ARP table mechanism
        self.m_arp_table = {}
        self.m_arp_table_update_event = threading.Event()
        self.m_arp_table_update_event.clear()
        self.m_arp_handler_lock = self.m_recv_pkt_thread.add_handler(self._arp_filter, self._arp_handler)
        
        logging.info("Network initialized")

    def _arp_filter(self, pkt):
        return isinstance(pkt.data, dpkt.arp.ARP)

    def _arp_handler(self, device, pkt):
        if 2 == pkt.data.op:
            # save sender MAC
            self.m_arp_table[pkt.data.spa] = pkt.data.sha
            self.m_arp_table_update_event.set()
        
        # always return True, to keep handler going
        return True
        
    def shutdown(self):
        self.m_recv_pkt_thread.kill()
    
    def resolve_ip(self, ip_addr):
        """resolve_ip(ip_addr) -> MAC address
        
        Resolve an IP address to MAC address using the cache or ARP
        """
        
        logging.debug("Resolving IP to MAC: %s", ip_addr)
        
        ip_addr_net = IP.ip_addr_aton(ip_addr)
        
        # check if address already in table
        if ip_addr_net in self.m_arp_table:
            logging.debug("MAC found in cache: %s", Ethernet.mac_addr_ntoa(self.m_arp_table[ip_addr_net]))
            return self.m_arp_table[ip_addr_net]
        
        # reset the update event
        self.m_arp_table_update_event.clear()
        
        logging.debug("Send ARP who-has broadcast")
        
        # send ARP request packet
        arp_pkt = ARP(self.m_src_mac, self.m_src_ip, ip_addr_net)
        self.send_eth_packet(ETHERNET_BROADCAST, arp_pkt, protocol=ETHERNET_TYPE_ARP)
        
        # wait for reply
        self.m_arp_table_update_event.wait(1000)
        
        # check if address already in table
        if ip_addr_net in self.m_arp_table:
            logging.debug("Reply stored in cache: %s", Ethernet.mac_addr_ntoa(self.m_arp_table[ip_addr_net]))
            return self.m_arp_table[ip_addr_net]
        
        raise Exception("No ARP reply after 1sec")
    
    def send_packet(self, pkt):
        """send_packet(pkt):
        
        Send a packet off the wire
        """
        
        return pkt.send(self.m_capture)
        
    def send_eth_packet(self, dst_mac, data, protocol=ETHERNET_TYPE_IP):
        """send_eth_packet(dst_mac, data)
        
        Send an Ethernet packet to a given destination address.
        """
        
        ether_pkt = Ethernet(self.m_src_mac, dst_mac, protocol, data)
        return self.send_packet(ether_pkt)
        
    def send_ip_packet(self, dst_addr, protocol, data):
        """send_ip_packet(dst_addr, protocol, data)
        
        Send an IP packet to a given destination IP address with a given protocol
        type. Destination MAC address is resolved automaticly.
        """
        
        ip_pkt = IP(self.m_src_ip, IP.ip_addr_aton(dst_addr), 100, protocol, data)
        dst_mac = self.resolve_ip(dst_addr)
        self.send_eth_packet(dst_mac, ip_pkt)
        
    def send_udp_packet(self, dst_addr, src_port, dst_port, data):
        """send_udp_packet(dst_addr, src_port, dst_port, data)
        
        Send a UDP packet to a given destination IP address with a specific
        source and destination port.
        """
        
        udp_packet = UDP(src_port, dst_port, data)
        self.send_ip_packet(dst_addr, UDP_PROTOCOL, udp_packet)

def open_first_device_for_capture():
    device_list = pywinpcap.find_all_devices()
    return pywinpcap.Capture(device_list[0]) 

def main():
    capture = open_first_device_for_capture()
    
    net = Network(capture)
    
    for x in range(100):
        net.send_udp_packet("192.168.1.1", 52, 53, "Test")
    
    try:
        while True:
            time.sleep(600)
    except KeyboardInterrupt:
        print "Goodbye"
    
    net.shutdown()

if __name__ == '__main__':
    main()